#include "kernel_operator.h"
using namespace AscendC;

// #define SCA
// #define PRINTF

#define BLOCK_BYTES_SIZE 32
#define BLOCK_STEP 1024

constexpr int32_t BUFFER_NUM = 2;

class HeavisideSpec1 {
   public:
    __aicore__ inline HeavisideSpec1() {}
    __aicore__ inline void Init(GM_ADDR input, GM_ADDR values, GM_ADDR out, uint32_t inputSize, uint32_t valueSize,
                                TPipe* pipeIn) {
        // printf("in spec1 init\n");

        this->coreNum = GetBlockNum();
        this->coreIndex = GetBlockIdx();
        this->pipe = pipeIn;
        this->inputSize = inputSize;
        this->valueSize = valueSize;
        this->MinBlockSize = 32 / sizeof(DTYPE_INPUT);
        this->BlockSize = BLOCK_BYTES_SIZE / sizeof(DTYPE_INPUT);
        this->BlockTotalNum = (inputSize + BlockSize - 1) / BlockSize;
        int BlockMod = BlockTotalNum % coreNum;

        this->BlockNum = (BlockTotalNum / coreNum) + (coreIndex < BlockMod ? 1 : 0);
        this->BlockBegin = ((BlockTotalNum / coreNum) * coreIndex) + (coreIndex < BlockMod ? coreIndex : BlockMod);
        // this->totalBlockSizeRow = ((BLOCK_BYTES_SIZE * dimSizeX + 255) / 256) * 256 / sizeof(DTYPE_Y);
        this->StepBlockSize = BlockSize * BLOCK_STEP;

        // printf("coreIndex coreNum BlockTotalNum BlockNum BlockBegin = %d %d %d %d %d \n" , coreIndex,coreNum,BlockTotalNum,BlockNum,BlockBegin);

        inputGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)input, inputSize * sizeof(DTYPE_INPUT));
        valuesGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)values, valueSize * sizeof(DTYPE_INPUT));
        outGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)out, inputSize * sizeof(DTYPE_INPUT));


        pipe->InitBuffer(InInput, BUFFER_NUM, BLOCK_STEP * BLOCK_BYTES_SIZE + 256);
        pipe->InitBuffer(InValues, BUFFER_NUM, 256);
        pipe->InitBuffer(OutOut, BUFFER_NUM, BLOCK_STEP * BLOCK_BYTES_SIZE + 256);

        pipe->InitBuffer(CmpBuf, (BLOCK_STEP * BLOCK_BYTES_SIZE ) / sizeof(DTYPE_INPUT) + 256);
    }

    __aicore__ inline void Process() {
        auto inValuesLocal = InValues.AllocTensor<DTYPE_INPUT>();
        DataCopy(inValuesLocal, valuesGm, BlockSize);
        InValues.EnQue(inValuesLocal);

        int end = (BlockBegin + BlockNum) * BlockSize;
        end = inputSize < end ? inputSize : end;


        int i = BlockBegin * BlockSize;
        int calNum = (end - i) < StepBlockSize ? (end - i) : StepBlockSize;

        // printf("i,calNum=%d %d\n",i,calNum);
        CopyIn(i, calNum);

        inValuesLocal = InValues.DeQue<DTYPE_INPUT>();
        SetValues = inValuesLocal.GetValue(0);

        // printf("SetValues=%f\n",SetValues);

        Compute(i, calNum);
        
        CopyOut(i, calNum);
        i += StepBlockSize;

        for (; i < end; i += StepBlockSize) {
            calNum = (end - i) < StepBlockSize ? (end - i) : StepBlockSize;
            // printf("i,calNum=%d %d\n",i,calNum);
            CopyIn(i, calNum);
            Compute(i, calNum);
            CopyOut(i, calNum);
        }

        InValues.FreeTensor(inValuesLocal);
    }

    __aicore__ inline void CopyIn(int begin, int calNum) {
        auto inInputLocal = InInput.AllocTensor<DTYPE_INPUT>();
        DataCopy(inInputLocal, inputGm[begin], calNum + MinBlockSize - 1);
        InInput.EnQue(inInputLocal);
    }

    __aicore__ inline void Compute(int begin, int calNum) {
        auto cmpBuf = CmpBuf.Get<uint8_t>();
        auto outOutLocal = OutOut.AllocTensor<DTYPE_INPUT>();
        Duplicate(outOutLocal, (DTYPE_INPUT)0, calNum);

        auto inInputLocal = InInput.DeQue<DTYPE_INPUT>();
        uint32_t calNumRow256 = ((calNum * sizeof(DTYPE_INPUT) + 255) / 256) * 256 / sizeof(DTYPE_INPUT);

        CompareScalar(cmpBuf, inInputLocal, (DTYPE_INPUT)0, CMPMODE::LE, calNumRow256);
        Select(outOutLocal, cmpBuf, outOutLocal, (DTYPE_INPUT)1, SELMODE::VSEL_TENSOR_SCALAR_MODE,
        calNumRow256);

        CompareScalar(cmpBuf, inInputLocal, (DTYPE_INPUT)0, CMPMODE::NE, calNumRow256);
        Select(outOutLocal, cmpBuf, outOutLocal, SetValues, SELMODE::VSEL_TENSOR_SCALAR_MODE,
        calNumRow256);

        OutOut.EnQue(outOutLocal);
        InInput.FreeTensor(inInputLocal);
    }

    __aicore__ inline void CopyOut(int begin, int calNum) {
        auto outOutLocal = OutOut.DeQue<DTYPE_INPUT>();
        DataCopy(outGm[begin], outOutLocal, calNum  + MinBlockSize - 1);
        OutOut.FreeTensor(outOutLocal);
    }

   private:
    TPipe* pipe;
    GlobalTensor<DTYPE_INPUT> inputGm;
    GlobalTensor<DTYPE_INPUT> valuesGm;
    GlobalTensor<DTYPE_INPUT> outGm;

    uint32_t BlockSize, BlockTotalNum, BlockNum, BlockBegin, totalBlockSizeRow, StepBlockSize, MinBlockSize;
    uint32_t coreIndex, coreNum;

    uint32_t inputSize, valueSize;
    DTYPE_INPUT SetValues;
    TBuf<QuePosition::VECCALC> CmpBuf;
    TQue<QuePosition::VECIN, BUFFER_NUM> InInput, InValues;
    TQue<QuePosition::VECOUT, BUFFER_NUM> OutOut;
};